#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@author: liang kang
@contact: gangkanli1219@gmail.com
@time: 1/12/18 1:30 PM
@desc: export a pd graph
"""
import os
import argparse

import tensorflow as tf


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--data-dir', type=str, default='', dest='data_dir',
        help='The directory where the checkpoint data is stored.')
    parser.add_argument(
        '--model-dir', type=str, default='', dest='model_dir',
        help='The directory where the pd model will be stored.')
    parser.add_argument(
        '--number', type=int, default=0, dest='number',
        help='The number of step to be trained.')
    return parser.parse_args()


def export(params):
    export_path_base = params.model_dir
    export_path = os.path.join(export_path_base, str(params.number))
    print('Exporting trained model to', export_path)
    if params.number != 0:
        model_pre = 'model.ckpt-{}'.format(params.number)
    else:
        model_pre = 'model.ckpt'
    with tf.Session() as sess:
        builder = tf.saved_model.builder.SavedModelBuilder(export_path)
        saver = tf.train.import_meta_graph(
            os.path.join(params.data_dir, '{}.meta'.format(model_pre)))
        saver.restore(sess, os.path.join(params.data_dir, model_pre))
        graph = tf.get_default_graph()
        image = graph.get_tensor_by_name('input_image:0')
        density_map = graph.get_tensor_by_name('density_map:0')
        category = graph.get_tensor_by_name('high_level:0')
        input_image = tf.saved_model.utils.build_tensor_info(image)
        output_map = tf.saved_model.utils.build_tensor_info(density_map)
        output_class = tf.saved_model.utils.build_tensor_info(category)
        inputs = {'images': input_image}
        outputs = {'density_map': output_map, 'class': output_class}
        signature = tf.saved_model.signature_def_utils.build_signature_def(
            inputs, outputs,
            tf.saved_model.signature_constants.PREDICT_METHOD_NAME)
        builder.add_meta_graph_and_variables(
            sess, ['frozen_model'],
            signature_def_map={'predict_images': signature})
        builder.save()

    print('Done exporting!')


if __name__ == '__main__':
    args = parse_args()
    export(args)
